"""
python fujian2_cnn_predict.py
"""

import pandas as pd
import numpy as np
from keras.models import Sequential
from keras.layers import Conv1D, MaxPooling1D, Flatten, Dense
import os
import json
import matplotlib.pyplot as plt


def predict_single_category(input_path, output_path):
    # 读取输入的 JSON 文件
    with open(input_path, 'r') as f:
        data = json.load(f)

    # 将 JSON 数据转换为 DataFrame
    df = pd.DataFrame(data)
    df['ds'] = pd.to_datetime(df['date'])
    df.rename(columns={'sales': 'y'}, inplace=True)

    # 准备数据用于训练
    values = df['y'].values.reshape(-1, 1)
    sequence_length = 10
    X = []
    y = []
    for i in range(len(values) - sequence_length):
        X.append(values[i:i + sequence_length])
        y.append(values[i + sequence_length])
    X = np.array(X)
    y = np.array(y)

    # 创建 CNN 模型
    model = Sequential()
    model.add(Conv1D(filters=64, kernel_size=2, activation='relu', input_shape=(sequence_length, 1)))
    model.add(MaxPooling1D(pool_size=2))
    model.add(Flatten())
    model.add(Dense(50, activation='relu'))
    model.add(Dense(1))

    # 编译模型
    model.compile(optimizer='adam', loss='mse')

    # 拟合模型
    model.fit(X, y, epochs=100, batch_size=32, verbose=0)

    # 进行预测
    last_sequence = values[-sequence_length:]
    predictions = []
    prediction_dates = []
    current_date = pd.to_datetime(df['ds'].iloc[-1]) + pd.Timedelta(days=1)
    for _ in range((pd.to_datetime('2023-09-30') - pd.to_datetime(df['ds'].iloc[-1])).days):
        new_prediction = model.predict(last_sequence.reshape(1, sequence_length, 1))
        predictions.append(new_prediction[0][0])
        prediction_dates.append(current_date)
        last_sequence = np.append(last_sequence[1:], new_prediction)
        current_date += pd.Timedelta(days=1)

    # 创建包含预测日期的数据框
    result_df = pd.DataFrame({'ds': prediction_dates, 'yhat': predictions})

    # 将结果转换回 JSON 格式
    result_json = result_df[['ds', 'yhat']].to_json(orient='records')
    result_json = json.loads(result_json)

    # 创建输出目录如果不存在
    output_dir = os.path.dirname(output_path)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 写入输出 JSON 文件
    with open(output_path, 'w') as f:
        json.dump(result_json, f)

    # 绘制原始数据和预测数据的图像
    plt.figure(figsize=(10, 6))
    plt.plot(df['ds'], df['y'], label='Original Data')
    plt.plot(result_df['ds'], result_df['yhat'], label='Predicted Data')
    plt.legend()

    # 保存图像
    plt.savefig(os.path.join(output_dir, os.path.basename(input_path).replace('.json', '.png')))


if __name__ == "__main__":
    input_path = "../fujian/fujian2/groupByCategory/category_category1.json"
    output_path = "./unit_test/category_category1_output.json"
    predict_single_category(input_path, output_path)
